from tensorflow.keras.layers import Input
from tensorflow.keras.callbacks import ModelCheckpoint
# from Tool.keras_callback import TensorBoard_eager_writer, confusion
from tensorflow import keras
import backbone.inception_v3 as inception
import backbone.BBN_net as BBN
from DataPipeline.Data import BNN_Data_tf14

# from Tool.tool import ConfusionMatrixMetrics
# tf.compat.v1.disable_eager_execution()

# tf.debugging.experimental.enable_dump_debug_info(dump_root="./out/logs/debug", tensor_debug_mode="FULL_HEALTH")
experment_name = "grad_tap"
data_path = "./test_data"
input_size = 512
batch_size = 2
model_param = {"blocks": [3, 4, 6, 3], "filter_lt": [64, 128, 256, 512], "radix": 2, "deep_stem": True,
               "avg_down": False, "using_basic_block": False}

# data_set = Customter(data_path, input_size)
data_set = BNN_Data_tf14(data_path, input_size=139)
train_data_set = data_set.train_data_set(batch_size)
val_data_set = data_set.val_data_set(batch_size)

x = Input((input_size, input_size, 3))
# base_model = resnet.create_model(x, data_set.class_num)
# base_model = resnest.create_model(x,data_set.class_num)
# base_model = inception.create_model(x, data_set.class_num)
# model = Estimate_tf14(base_model,keras.layers.Input(shape=(4,)))
model = BBN.create_model(batch_size, input_size=139, classes=data_set.class_num)
# model = BBN.load_model("./out/temp_model-val_acc_0.375.h5")
model.summary()
# model_path = "out/" + experment_name + "_9_{epoch:03d}-{val_loss:.4f}"
model_path = "out/" + experment_name
checkpoint_callback = ModelCheckpoint(model_path, monitor='val_loss', verbose=1, save_best_only=False,
                                      save_weights_only=False, mode='auto', period=1)
# tensorboard = TensorBoard_eager_writer("./out/logs", experment_name, histogram_freq=5, write_graph=True,
#                                        write_images=False, update_freq='batch', profile_batch=2, embeddings_freq=0)
# confusion_call = confusion(tensorboard.train_writer, tensorboard.val_writer, data_set.class_names)
# model = Estimate(inputs=base_model.input, outputs=base_model.output,
#                 train_summary_writer=tensorboard.train_writer,
#                  # test_summary_writer=tensorboard.val_writer
#               )
# model.compile(tf.keras.optimizers.SGD(), tf.keras.losses.categorical_crossentropy,
#               ["accuracy", ConfusionMatrixMetrics(data_set.class_num)])
model.fit([train_data_set,train_data_set], epochs=2, steps_per_epoch=data_set.train_num // batch_size,
          validation_data=val_data_set, validation_steps=data_set.test_num // batch_size,
          callbacks=[checkpoint_callback
                     # tensorboard, confusion_call
                     ]
          )

# model.load_weights("./out/grad_tap_9_001-0.3750/variables/variables")
# estimater.fit(x=train_data_set, epochs=1, steps_per_epoch=data_set.train_num // batch_size,
#               validation_data=val_data_set, validation_steps=data_set.test_num // batch_size,
#               callbacks=[checkpoint_callback, tensorboard]
#               )

# from Tool.Visualization import grad_cam_image
# grad_cam_image(model,"./test_data/1/201912230001_20200421001.jpg")
